import os
from typing import Dict
import numpy as np

# from utils.log import get_logger


def write_results(filename, results, data_type):
    if data_type == 'mot':
        save_format = '{frame},{id},{x1},{y1},{w},{h},-1,-1,-1,-1\n'
    elif data_type == 'kitti':
        save_format = '{frame} {id} pedestrian 0 0 -10 {x1} {y1} {x2} {y2} -10 -10 -10 -1000 -1000 -1000 -10\n'
    else:
        raise ValueError(data_type)

    with open(filename, 'w') as f:
        for frame_id, tlwhs, track_ids in results:
            if data_type == 'kitti':
                frame_id -= 1
            for tlwh, track_id in zip(tlwhs, track_ids):
                if track_id < 0:
                    continue
                x1, y1, w, h = tlwh
                x2, y2 = x1 + w, y1 + h
                line = save_format.format(frame=frame_id, id=track_id, x1=x1, y1=y1, x2=x2, y2=y2, w=w, h=h)
                f.write(line)


# def write_results(filename, results_dict: Dict, data_type: str):
#     if not filename:
#         return
#     path = os.path.dirname(filename)
#     if not os.path.exists(path):
#         os.makedirs(path)

#     if data_type in ('mot', 'mcmot', 'lab'):
#         save_format = '{frame},{id},{x1},{y1},{w},{h},1,-1,-1,-1\n'
#     elif data_type == 'kitti':
#         save_format = '{frame} {id} pedestrian -1 -1 -10 {x1} {y1} {x2} {y2} -1 -1 -1 -1000 -1000 -1000 -10 {score}\n'
#     else:
#         raise ValueError(data_type)

#     with open(filename, 'w') as f:
#         for frame_id, frame_data in results_dict.items():
#             if data_type == 'kitti':
#                 frame_id -= 1
#             for tlwh, track_id in frame_data:
#                 if track_id < 0:
#                     continue
#                 x1, y1, w, h = tlwh
#                 x2, y2 = x1 + w, y1 + h
#                 line = save_format.format(frame=frame_id, id=track_id, x1=x1, y1=y1, x2=x2, y2=y2, w=w, h=h, score=1.0)
#                 f.write(line)
#     logger.info('Save results to {}'.format(filename))


def read_results(filename, data_type: str, is_gt=False, is_ignore=False):
    if data_type in ('mot', 'lab'):
        read_fun = read_mot_results
    else:
        raise ValueError('Unknown data type: {}'.format(data_type))

    return read_fun(filename, is_gt, is_ignore)



def read_mot_results(filename, is_gt, is_ignore):
    """
    读取MOT数据格式
    
    dict[frame_id -> (tlwh_bbox_tuple , id , score)]
    
    Param
    ---
    * `filename` : MOT格式的文件路径
    * `is_gt` : 读取的文件是否是gt文件,优先于`is_ignore`,自动忽略`mark == 0`与`label not in valid_labels`的条目
    * `is_ignore` : 读取忽略的类别标签,提取`label not in ignore_labels`且`vis_ratio >= 0`的条目

    gt.txt
    ---
    * `is_gt = True` 或 `is_ignore = True`,`is_gt`优先于`is_ignore`
    * 9个参数 : 0<frame_id>,1<track_id>,2<bb_left>,3<bb_top>,4<bb_width>, \
        5<bb_height>,6<mark>,7<label>,8<vis_ratio>
    * 自动忽略`mark == 0`与`label not in valid_labels`的条目
    
    det.txt
    ---
    * `is_gt = False` 且 `is_ignore = False`
    * 10个参数 : 0<frame_id>,1<id>,2<bb_left>,3<bb_top>,4<bb_width>, \
        5<bb_height>,6<conf>,-1,-1,-1
    * 提取`label not in ignore_labels`且`vis_ratio >= 0`的条目
    """
    valid_labels = {}
    ignore_labels = {}
    results_dict = dict()
    if os.path.isfile(filename):
        with open(filename, 'r') as f:
            for line in f.readlines():
                linelist = line.split(',')
                if len(linelist) < 7:
                    continue
                # 获得frame编号
                fid = int(linelist[0])
                if fid < 1:
                    continue
                # 如果frame没在dict内,需要初始化为list()
                # 如果键值中存在frame,则不需要初始化
                results_dict.setdefault(fid, list())

                if is_gt: # 是gt文件,忽略指定的轨迹与特定的类别
                    # if 'MOT16-' in filename or 'MOT17-' in filename:
                    label = int(float(linelist[7])) # 目标的种类ID
                    mark = int(float(linelist[6])) # 目标轨迹是否进入考虑范围内的标志
                    # mark == 0:目标轨迹不在考虑范围内
                    # label not in valid_labels: 目标轨迹不是需要考虑的目标种类
                    if mark == 0 or label not in valid_labels:
                        continue
                    score = 1
                elif is_ignore: # gt文件
                    # if 'MOT16-' in filename or 'MOT17-' in filename:
                    label = int(float(linelist[7])) # 目标的种类ID
                    vis_ratio = float(linelist[8]) # 目标被遮挡时的可见比率
                    if label not in ignore_labels and vis_ratio >= 0:
                        continue
                    score = 1
                else: # det文件
                    # 结果文件的每行的第7个变量必须是score
                    score = float(linelist[6])

                # 将bbox的4个坐标转换成tuple形式
                tlwh = tuple(map(float, linelist[2:6]))
                # 获取追踪目标的ID
                target_id = int(linelist[1])

                # frame_id -> ( tlwh_bbox_tuple , id , score )
                results_dict[fid].append((tlwh, target_id, score))

    return results_dict


def unzip_objs(objs):
    if len(objs) > 0:
        tlwhs, ids, scores = zip(*objs)
    else:
        tlwhs, ids, scores = [], [], []
    tlwhs = np.asarray(tlwhs, dtype=float).reshape(-1, 4)

    return tlwhs, ids, scores